#include "tiling/platform/platform_ascendc.h"
#include "mish_grad_tiling.h"
#include "register/op_def_registry.h"

namespace optiling
{
    const uint32_t BLOCK_SIZE = 32;
    static ge::graphStatus TilingFunc(gert::TilingContext *context)
    {
        MishGradTilingData tiling;

        // get ubsize
        uint64_t ubSize;
        auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
        ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);

        // datatype
        uint32_t sizeofdatatype;
        auto dt = context->GetInputTensor(0)->GetDataType();
        if (dt == 1)
        { // float16
            sizeofdatatype = 2;
        }
        else if (dt == 0)
        { // float
            sizeofdatatype = 4;
        }

        // totalNum, totalLength
        uint32_t grad_size = context->GetInputShape(0)->GetStorageShape().GetShapeSize();
        uint32_t x_size = context->GetInputShape(1)->GetStorageShape().GetShapeSize();
        auto tanhx = context->GetInputTensor(2);
        bool tanhx_opt = (tanhx) ? true : false;
        bool broadcast = (grad_size == x_size) ? false : true;
        uint32_t totalNum = (grad_size > x_size) ? grad_size : x_size; // 输入总元素个数
        uint32_t totalLength = totalNum * sizeofdatatype;              // 输入总bit

        // ub
        uint32_t ubDataNumber = 8;
        uint32_t ub_block_num = ((ubSize) / BLOCK_SIZE / sizeofdatatype) / ubDataNumber; // 一次ub搬运最多操作多少个block
        uint32_t ub_data_num = (ub_block_num * BLOCK_SIZE) / sizeofdatatype;

        // totalLengthAligned
        uint32_t totalLengthAligned;
        if (totalLength % BLOCK_SIZE != 0)
        {
            totalLengthAligned =
                ((totalLength + BLOCK_SIZE - 1) / BLOCK_SIZE) * BLOCK_SIZE;
        }
        else
        {
            totalLengthAligned = totalLength;
        }

        // everyCoreBlockNum
        uint32_t Block_num = totalLengthAligned / BLOCK_SIZE; // 需要多少个Block(32bit)
        uint32_t totalNumAligned = totalLengthAligned / sizeofdatatype;
        // everyCoreInputDataNum
        uint32_t CoreDataNum = Block_num * BLOCK_SIZE / sizeofdatatype; // 每个核处理的输入元素个数

        // tile
        uint32_t tile_num = Block_num / ub_block_num;
        uint32_t final_tile_num = (Block_num % ub_block_num) == 0 ? tile_num : tile_num + 1;

        // lasttileNum
        uint32_t tailNum = CoreDataNum - (tile_num * ub_data_num);
        uint32_t lasttileNum = (tailNum == 0) ? ub_data_num : tailNum;
        context->SetBlockDim(1);

        printf("totalNum: %d \n", totalNumAligned);
        printf("tileNum: %d \n", final_tile_num);
        printf("tileLength: %d \n", ub_data_num);
        printf("lasttileNum: %d \n", lasttileNum);
        if (tanhx_opt)
        {
            printf("tanhx isn't None\n");
        }
        else
        {
            printf("tanhx is None\n");
        }
        tiling.set_totalNum(totalNum);
        tiling.set_tileNum(final_tile_num);
        tiling.set_tileLength(ub_data_num);
        tiling.set_lasttileNum(lasttileNum);
        tiling.set_tanhx_opt(tanhx_opt);
        tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
        context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
        return ge::GRAPH_SUCCESS;
    }
}

namespace ge
{
    static ge::graphStatus InferShape(gert::InferShapeContext *context)
    {
        const gert::Shape *x1_shape = context->GetInputShape(0);
        gert::Shape *y_shape = context->GetOutputShape(0);
        *y_shape = *x1_shape;
        return GRAPH_SUCCESS;
    }
}

namespace ops
{
    class MishGrad : public OpDef
    {
    public:
        explicit MishGrad(const char *name) : OpDef(name)
        {
            this->Input("grad")
                .ParamType(REQUIRED)
                .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
                .Format({ge::FORMAT_ND, ge::FORMAT_ND})
                .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
            this->Input("x")
                .ParamType(REQUIRED)
                .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
                .Format({ge::FORMAT_ND, ge::FORMAT_ND})
                .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
            this->Input("tanhx")
                .ParamType(OPTIONAL)
                .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
                .Format({ge::FORMAT_ND, ge::FORMAT_ND})
                .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
            this->Output("y")
                .ParamType(REQUIRED)
                .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
                .Format({ge::FORMAT_ND, ge::FORMAT_ND})
                .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});

            this->SetInferShape(ge::InferShape);

            this->AICore()
                .SetTiling(optiling::TilingFunc);
            this->AICore().AddConfig("ascend310b");
        }
    };

    OP_ADD(MishGrad);
}
